import unittest

from psychsim.pwl import *
from psychsim.action import Action,ActionSet
from psychsim.world import World,stateKey
from psychsim.agent import Agent
from psychsim.reward import *

'''
This example examines belief updates over stochastic dynamics.
'''

class FairCoinTestCase(unittest.TestCase):

    def setUp(self):
        #setup world, agent, and basics
        self.world = World()
        self.obs = Agent('observer')
        self.world.addAgent(self.obs)
        self.coin = Agent('coin')
        self.world.addAgent(self.coin)
        self.world.setOrder([self.coin.name, self.obs.name])

    def addStates(self):
        #define states
        self.world.defineState(None, 'heads', int, lo=0, hi=10, description='Number of heads.')
        self.world.setState(None,'heads', 0)

        self.world.defineState(self.coin.name, 'head prob', float, lo=0.0, hi=1.0, description='Probability of heads.')
        self.world.setState(self.coin.name, 'head prob', 1.0)

        self.world.defineState(None, 'round', int, lo=1, hi=10, description='The current round.')
        self.world.setState(None, 'round', 1)
                    
    def addTermination(self):
        #terminate if maxRounds reached
        self.world.addTermination(makeTree({'if': equalRow(stateKey(None, 'round'), 10), True: True, False: False}))

    def addActions(self):
        #add actions and appropriate dynamics
        self.obs.addAction({'verb': 'do nothing'})
        self.coin.addAction({'verb': 'flip'})

        #set legality
        for action in filterActions({'verb': 'do nothing'}, self.obs.actions):
            self.obs.setLegal(action, makeTree({'if': thresholdRow(stateKey(None, 'round'), 10),
                                            True: False, False: True}))

        #define dynamics
        for action in filterActions({'verb': 'flip'}, self.coin.actions):
            self.coin.setLegal(action, makeTree({'if': thresholdRow(stateKey(None, 'round'), 10),
                                            True: False, False: True}))
            prob = stateKey(self.coin, 'head prob')
            heads = stateKey(None, 'heads')
            tree = makeTree({'if': thresholdRow(prob, 0.99),
                             True: {'distribution': [(incrementMatrix(heads, 1), 1.0),
                                                     (incrementMatrix(heads, 0), 0.0)]},
                             False: {'distribution': [(incrementMatrix(heads, 1), 0.5),
                                                      (incrementMatrix(heads, 0), 0.5)]}})
            self.world.setDynamics(stateKey(None, 'heads'), action, tree)

            tree = makeTree(incrementMatrix(stateKey(None, 'round'), 1))
            self.world.setDynamics(stateKey(None, 'round'), action, tree)

    def setModelAndRewards(self):
        #set models, rewards, and beliefs
        key = stateKey(self.coin.name, 'head prob')

        self.coin.addModel('fair',rationality=10., selection='distribution', parent=True)
        self.coin.setBelief(key, Distribution({0.5: 1.0}), 'fair')

        goalMoreHeads = maximizeFeature(stateKey(None,'heads'))        
        self.coin.addModel('rigged', R={goalMoreHeads: 1.0},rationality=10., selection='distribution', parent=True)
        self.coin.setBelief(key, Distribution({1.0: 1.0, 0.0: 0.0}), 'rigged')

        self.world.setMentalModel(self.obs.name, self.coin.name, {'fair': 0.5, 'rigged': 0.5})
        self.world.setModel(self.obs.name, True)

    def setAttributes(self):
        #set miscellaneous attributes 
        self.obs.setAttribute('horizon', 3)

    def testSetup(self):
        self.addStates()
        self.addTermination()
        self.addActions()
        self.setModelAndRewards()

    def testRun(self):
        self.testSetup()
        self.world.printState(beliefs=True)
        self.world.state.normalize()
        while True: #should probably check for correctness here, which is messy with stochastic dynamics
            self.world.explain(self.world.step(),level=3)
            self.world.state.select()
            self.world.printState(beliefs=True)
            self.world.printModels()
            if self.world.terminated():
                break

if __name__ == "__main__":
    suite = unittest.TestSuite()
    for method in dir(FairCoinTestCase):
       if method.startswith("test"):
          suite.addTest(FairCoinTestCase(method))
    unittest.TextTestRunner().run(suite)
            